import numpy as np
import argparse

def sishewuru(x: float) -> int:
    if x > 0.0: return int(x + 0.5)
    elif x == 0.0: return 0
    elif x < 0.0: return int(x - 0.5)
    return 0

def initDNA():
    pa, pb, pc = 0.99, 0.005, 0.0025
    PAM_DNA = np.array([[pa, pb, pc, pc], [pb, pa, pc, pc], [pc, pc, pa, pb], [pc, pc, pb, pa]], dtype = np.float64)
    order = "ATCG"
    S = PAM_DNA
    for i in range(199): S = np.matmul(S, PAM_DNA)
    # print(S)
    S = 1000 * np.log10(4. * S)
    # print(S)
    ave = sum([sum(i / 16.0 for i in j) for j in S])
    #print(ave)
    S = np.array([[i - ave for i in j] for j in S])
    # print(S)
    ave = 0.0
    for i in range(0, 4): 
        ave += S[i][i] / 4.
    S = np.array([[i * 600.0 / ave for i in j] for j in S])
    #print(S)
    S = np.array([[sishewuru(i) for i in j] for j in S])
    print("The Score Matrix of DNA is: \n", end = '     ')
    for i in range(4):
        print(order[i], end = '    ')
    print()
    for i in range(4):
        print(order[i], end = '')
        for j in range(4):
            print('{0:5d}'.format(S[i][j]), end = '')
        print()

def initProtein(rescale = True, offset: float = 0.0) -> None:
    '''
    '''
    BLOSUM62 = [[5.89,-2.12,-2.30,-2.63,-0.61,-1.21,-1.30,0.24,-2.44,-1.98,-2.20,-1.10,-1.40,-3.32,-1.22,1.67,-0.07,-3.79,-2.65,-0.28], 
                [-2.12,8.21,-0.66,-2.41,-5.08,1.47,-0.17,-3.46,-0.37,-4.49,-3.23,3.16,-2.05,-4.18,-3.16,-1.15,-1.68,-4.02,-2.54,-3.75], 
                [-2.30,-0.66,8.48,1.91,-3.99,0.00,-0.40,-0.63,0.87,-4.83,-5.07,-0.27,-3.23,-4.49,-3.00,0.90,-0.07,-5.54,-3.12,-4.31], 
                [-2.63,-2.41,1.91,8.66,-5.19,-0.47,2.27,-1.97,-1.68,-4.68,-5.41,-1.05,-4.59,-5.23,-2.22,-0.39,-1.58,-6.32,-4.60,-4.71], 
                [-0.61,-5.08,-3.99,-5.19,12.87,-4.35,-5.42,-3.75,-4.48,-1.84,-1.92,-4.55,-2.13,-3.56,-4.19,-1.31,-1.30,-3.46,-3.61,-1.21], 
                [-1.21,1.47,0.00,-0.47,-4.35,7.93,2.78,-2.68,0.67,-4.15,-3.20,1.91,-0.63,-4.75,-1.92,-0.15,-1.01,-2.92,-2.13,-3.30], 
                [-1.30,-0.17,-0.40,2.27,-5.42,2.78,7.35,-3.17,-0.18,-4.79,-4.27,1.16,-3.00,-4.79,-1.67,-0.22,-1.29,-4.25,-3.03,-3.66], 
                [0.24,-3.46,-0.63,-1.97,-3.75,-2.68,-3.17,8.34,-3.06,-5.59,-5.44,-2.29,-4.01,-4.66,-3.20,-0.44,-2.36,-3.74,-4.56,-4.71], 
                [-2.44,-0.37,0.87,-1.68,-4.48,0.67,-0.18,-3.06,11.27,-4.85,-4.18,-1.08,-2.33,-1.85,-3.24,-1.32,-2.53,-3.51,2.54,-4.68], 
                [-1.98,-4.49,-4.83,-4.68,-1.84,-4.15,-4.79,-5.59,-4.85,6.00,2.28,-4.01,1.69,-0.24,-4.14,-3.52,-1.08,-3.87,-2.00,3.82], 
                [-2.20,-3.23,-5.07,-5.41,-1.92,-3.20,-4.27,-5.44,-4.18,2.28,5.77,-3.67,2.99,0.62,-4.29,-3.66,-1.80,-2.45,-1.59,1.18], 
                [-1.10,3.16,-0.27,-1.05,-4.55,1.91,1.16,-2.29,-1.08,-4.01,-3.67,6.76,-2.03,-4.62,-1.52,-0.31,-1.00,-4.43,-2.73,-3.39], 
                [-1.40,-2.05,-3.23,-4.59,-2.13,-0.63,-3.00,-4.01,-2.33,1.69,2.99,-2.03,8.09,0.02,-3.71,-2.22,-1.00,-2.14,-1.49,1.03], 
                [-3.32,-4.18,-4.49,-5.23,-3.56,-4.75,-4.79,-4.66,-1.85,-0.24,0.62,-4.62,0.02,9.07,-5.40,-3.55,-3.16,1.38,4.41,-1.27], 
                [-1.22,-3.16,-3.00,-2.22,-4.19,-1.92,-1.67,-3.20,-3.24,-4.14,-4.29,-1.52,-3.71,-5.40,11.05,-1.21,-1.61,-5.48,-4.38,-3.52], 
                [1.67,-1.15,0.90,-0.39,-1.31,-0.15,-0.22,-0.44,-1.32,-3.52,-3.66,-0.31,-2.22,-3.55,-1.21,5.83,2.07,-4.13,-2.53,-2.47], 
                [-0.07,-1.68,-0.07,-1.58,-1.30,-1.01,-1.29,-2.36,-2.53,-1.08,-1.80,-1.00,-1.00,-3.16,-1.61,2.07,6.82,-3.64,-2.41,-0.08], 
                [-3.79,-4.02,-5.54,-6.32,-3.46,-2.92,-4.25,-3.74,-3.51,-3.87,-2.45,-4.43,-2.14,1.38,-5.48,-4.13,-3.64,15.76,3.23,-4.25], 
                [-2.65,-2.54,-3.12,-4.60,-3.61,-2.13,-3.03,-4.56,2.54,-2.00,-1.59,-2.73,-1.49,4.41,-4.38,-2.53,-2.41,3.23,9.89,-1.81], 
                [-0.28,-3.75,-4.31,-4.71,-1.21,-3.30,-3.66,-4.71,-4.68,3.82,1.18,-3.39,1.03,-1.27,-3.52,-2.47,-0.08,-4.25,-1.81,5.65]]
                         
    order = "ARNDCQEGHILKMFPSTWYV"
    '''
    if len(BLOSUM62) != 20: raise TypeError("Protein must be 20 types")
    for i in BLOSUM62:
        if len(i) != 20: raise TypeError("Protein must be 20 types")
    '''
    freq = [0.077, 0.051, 0.043, 0.052, 0.02, 0.041, 0.062, 0.074, 0.023, 0.052, 0.091, 0.059, 0.024, 0.04, 0.051, 0.069, 0.059, 0.014, 0.032, 0.066]
    ave = 0.0
    for i in range(len(BLOSUM62)):
        for j in range(len(BLOSUM62[i])):
            ave += BLOSUM62[i][j] * freq[i] * freq[j]
    # print('ave2', ave) # -1.5026337900000002
    if rescale: 
        BLOSUM62 = [[i - ave for i in j] for j in BLOSUM62]
        #print(BLOSUM62)
    ave = 0.0
    for i in range(len(BLOSUM62)):
        ave += BLOSUM62[i][i] * freq[i]
    # print('ave1', ave) # 9.175523789999998
    if rescale:
        BLOSUM62 = [[i * 600.0 / ave for i in j] for j in BLOSUM62]
        #print(BLOSUM62)
    else:
        BLOSUM62 = [[i * 600.0 for i in j] for j in BLOSUM62]
    if offset != 0.0:
        BLOSUM62 = [[i - offset for i in j] for j in BLOSUM62]
    BLOSUM62 = [[sishewuru(i) for i in j] for j in BLOSUM62]
    print("char order[20] = {", end = '')
    for i in range(20):
        if i != 19: print("'" + order[i] + "', ", end = '')
        else: print("'" + order[i] + "'", end = '')
    print("};\nint BLOSUM62[20][20] = {\n{", end = '')
    for i in range(20):
        for j in range(20):
            if j != 19: print('{0:5d}'.format(BLOSUM62[i][j]) + ",", end = '')
            else: print('{0:5d}'.format(BLOSUM62[i][j]), end = '')
        if i != 19: print("}, \n{", end = '')
        else: print("}\n};")
            

def argparser():
    '''
    define a parser to process file name\n
    return: processed arguments
    '''
    parser = argparse.ArgumentParser(description = "Initial the DNA and Protein scoring matrix. ")
    parser.add_argument("--type", required = True, default = None, type = str, choices = ["DNA", "Protein"], help = "The type of matrix, MUST be in ['DNA', 'Protein']")
    return parser.parse_args()


if __name__ == '__main__':
    Type = argparser().type
    if Type == "DNA":
        initDNA()
    else: 
        initProtein()